#!/usr/bin/perl

use strict;
use utf8;
use XML::Parser;
use Getopt::Std;

binmode(STDIN, ':encoding(utf8)');
binmode(STDOUT, ':encoding(utf8)');
binmode(STDERR, ':encoding(utf8)');

my %options;
my %dict;
my %dict_w;
my %d;
my $total_docs = 0;
my %bigrams;
my $curr_token_id = undef;

# globals for tag_start / tag_end
my $prev_word = undef;

getopts('Ccf:iln:opsb:', \%options);
# f - path to xml with annotation

if (!$options{'f'}) {
    print STDERR "-f argument not specified, presume standard input as file\n";
}

if ($options{'b'}) {
    load_bigrams($options{'b'}, 5);
}

# xml parsing

my $parser = XML::Parser->new(Handlers=> {Start => \&tag_start, End => \&tag_end});
if ($options{'f'}) {
    $parser->parsefile($options{'f'});
} else {
    $parser->parse(*STDIN);
}

# output
my $total_keys = 0;
for my $k(sort keys %dict) {
    printf "%s\t%d\t%d\n", $k, $dict_w{$k}, $dict{$k} ;
    $total_keys += $dict_w{$k};
}

print "###TOTAL###\t$total_keys\t$total_docs\n";

# subroutines

sub tag_start {
    my ($expat, $tag_name, %attr) = @_;
    
    if ($tag_name eq 'token' || ($options{'l'} && $tag_name eq 'l')) {
        my $tt = undef;

        if ($tag_name eq 'token') {
          $curr_token_id = $attr{'id'};
          if ($options{'l'}) {
            return;
          }

          if ($attr{'text'} !~ /[А-ЯЁа-яё0-9]/) {
            return;
          }

          $tt = to_lower($attr{'text'});
        } elsif ($tag_name eq 'l' && defined($curr_token_id)) {
          if ($attr{'t'} !~ /[А-ЯЁа-яё0-9]/) {
            return;
          }

          $tt = to_lower($attr{'t'});
          $curr_token_id = undef;
        }

        $tt =~ tr/[Ёё]/[Ее]/;
        $d{$tt}++;
        $dict_w{$tt}++;

        # calculate bigrams
        if (defined($prev_word)) {
          my $k = $prev_word . "_" . $tt;
          if (exists($bigrams{$k})) {
            $d{$k}++;
            $dict_w{$k}++;
          }
        }

        $prev_word = $tt;
    }
}
sub tag_end {
    my ($expat, $tag_name) = @_;

    if ($tag_name eq 'text') {
        foreach my $k (keys %d) {
            $dict{$k} += 1;
        }
        %d = ();
        $total_docs += 1;
        $prev_word = undef;
    } elsif ($tag_name eq 'token') {
        $curr_token_id = undef;
    }
}

sub load_bigrams {
  my ($fn, $limit) = @_;
  open(F, "< $fn");
  binmode(F, ":encoding(utf-8)");
  while (<F>) {
    chomp $_;
    my ($w1, $w2, $n1, $n2, $nbi, $mi) = split(/\s+/, $_);
    if ($mi < $limit) {
      last;
    }
    my $k = $w1 . "_" . $w2;
    $bigrams{$k} = $mi;
  }
  close(F);
#  print STDERR scalar(keys %bigrams) . " bigrams loaded\n";
}

sub to_lower {
    my $s = lc(shift);
    $s =~ tr/А-ЯЁ/а-яё/;
    return $s;
}

sub to_upper {
    my $s = lc(shift);
    $s =~ tr/а-яё/А-ЯЁ/;
    return $s;
}
